use crate::hot100::base::{ListNode, Vec2ListNode};

#[test]
fn test_merge_k_lists() {
    let node = ListNode::new(0);
    assert_eq!(
        merge_k_lists(vec![
            node.from_vec(vec![1, 4, 5]),
            node.from_vec(vec![1, 3, 4]),
            node.from_vec(vec![2, 6])
        ]),
        node.from_vec(vec![1, 1, 2, 3, 4, 4, 5, 6])
    );

    assert_eq!(merge_k_lists(vec![]), node.from_vec(vec![]));

    assert_eq!(
        merge_k_lists(vec![node.from_vec(vec![])]),
        node.from_vec(vec![])
    );
}

/*
   23. 合并K个升序链表
   给你一个链表数组，每个链表都已经按升序排列。

   请你将所有链表合并到一个升序链表中，返回合并后的链表。

   示例 1：
   输入：lists = [[1,4,5],[1,3,4],[2,6]]
   输出：[1,1,2,3,4,4,5,6]
   解释：链表数组如下：
   [
       1->4->5,
       1->3->4,
       2->6
   ]
   将它们合并到一个有序链表中得到。
   1->1->2->3->4->4->5->6

   示例 2：
   输入：lists = []
   输出：[]

   示例 3：
   输入：lists = [[]]
   输出：[]

   提示：
   k == lists.length
   0 <= k <= 10^4
   0 <= lists[i].length <= 500
   -10^4 <= lists[i][j] <= 10^4
   lists[i] 按 升序 排列
   lists[i].length 的总和不超过 10^4
*/
pub fn merge_k_lists(lists: Vec<Option<Box<ListNode>>>) -> Option<Box<ListNode>> {
    if lists.is_empty() {
        return None;
    }

    let mut head = lists;
    let (left, right) = (0, head.len() - 1);
    merge(&mut head, left, right)
}

fn merge(
    lists: &mut Vec<Option<Box<ListNode>>>,
    left: usize,
    right: usize,
) -> Option<Box<ListNode>> {
    if left == right {
        return lists[left].take();
    }

    let mid = left + (right - left) / 2;
    let l1 = merge(lists, left, mid);
    let l2 = merge(lists, mid + 1, right);
    merge_two_lists(l1, l2)
}

fn merge_two_lists(l1: Option<Box<ListNode>>, l2: Option<Box<ListNode>>) -> Option<Box<ListNode>> {
    let mut dummy = ListNode::new(0);
    let (mut l1, mut l2, mut p) = (l1, l2, &mut dummy);

    while let (Some(n1), Some(n2)) = (l1.as_ref(), l2.as_ref()) {
        if n1.val < n2.val {
            p.next = l1;
            p = p.next.as_mut()?;
            l1 = p.next.take();
        } else {
            p.next = l2;
            p = p.next.as_mut()?;
            l2 = p.next.take();
        }
    }
    p.next = l1.or(l2);

    dummy.next
}
